import numpy as np
import matplotlib.pyplot as plt
import math
import mpl_toolkits.axisartist as axisartist


def triangle_signal(f):
    if f < -50 or f > 50:
        return 0
    elif f < 0:
        return f / 50 + 1
    else:
        return 1 - f / 50

def ifft(t, delta):
    return math.sin(math.pi * delta * (-t)) ** 2 / (math.pi ** 2 * delta * (-t) ** 2)

def sample_convet(f, fs):
    amplitude = 0
    for i in range(-10, 11):
        amplitude += fs * triangle_signal(f - i * fs)
    return amplitude

def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)


if __name__ == "__main__":

    fig = plt.figure(figsize=(12, 8))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    f = np.arange(-50, 50, 0.01)
    T = 100
    delta = 50

    freq_signal = np.array([triangle_signal(i) for i in f])
    set_axis(fig, 221, "带限信号频率域图像", (-50, 50), (-1, 3))
    plt.plot(f, freq_signal)

    t = np.arange(-10, 10, 0.001)
    time_signal = np.array([ifft(i, delta) for i in t])
    set_axis(fig, 222, "带限信号时间域图像", (-0.1, 0.1), (-2, 50))
    plt.plot(t, time_signal)

    fs = 50
    f = np.arange(-150, 150, 0.01)
    freq_signal = np.array([sample_convet(i, fs) for i in f])
    set_axis(fig, 234, "50Hz采样后信号的频谱", (-150, 150), (-1, 60))
    plt.plot(f, freq_signal)
    x = np.arange(-50, 50, 0.01)
    y = np.array([50 * triangle_signal(i) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(0, 100, 0.01)
    y = np.array([50 * triangle_signal(i-50) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(50, 150, 0.01)
    y = np.array([50 * triangle_signal(i-100) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(100, 150, 0.01)
    y = np.array([50 * triangle_signal(i-150) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(-100, 0, 0.01)
    y = np.array([50 * triangle_signal(i+50) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(-150, -50, 0.01)
    y = np.array([50 * triangle_signal(i+100) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    x = np.arange(-150, -100, 0.01)
    y = np.array([50 * triangle_signal(i+150) for i in x])
    plt.plot(x, y, color='red', linestyle=':')
    

    fs = 100
    f = np.arange(-150, 150, 0.01)
    freq_signal = np.array([sample_convet(i, fs) for i in f])
    set_axis(fig, 235, "100Hz采样后信号的频谱", (-150, 150), (-1, 110))
    plt.plot(f, freq_signal)

    fs = 200
    f = np.arange(-300, 300, 0.01)
    freq_signal = np.array([sample_convet(i, fs) for i in f])
    set_axis(fig, 236, "100Hz采样后信号的频谱", (-300, 300), (-1, 210))
    plt.plot(f, freq_signal)

    plt.show()